import mxnet.ndarray as nd
import numpy as np
import torch
class TorchWrapper():
    def __init__(self, mx_metric):
        self.mx_metric = mx_metric
    def update(self, labels, preds):
        if isinstance(labels, torch.Tensor):
            labels = labels.data.cpu().numpy()
        if isinstance(preds, torch.Tensor):
            preds  = preds.data.cpu().numpy()
        if labels is not None:
            labels = nd.array([labels] if labels.shape is () else labels)
        if preds is not None:
            preds = nd.array([preds] if preds.shape is () else preds)
        self.mx_metric.update(labels,preds)
    def get(self):
        return self.mx_metric.get()
    def reset(self):
        self.mx_metric.reset()
